using UnityEngine;
using UnityEngine.Profiling;
using UnityEngine.Rendering;
using System.Collections.Generic;

public class IndirectEXP: MonoBehaviour
{
    public bool useFrustumCull = true;
    public bool useHiZCull = true;
    public int instanceCount = 10000;
    public HiZBuffer m_hizBuffer;
    public Mesh instanceMesh;
    public Material instanceMaterial;
    public Camera mainCamera;
    public int subMeshIndex = 0;

    private int cachedInstanceCount = -1;
    private int cachedSubMeshIndex = -1;
    private uint[] args = new uint[5] { 0, 0, 0, 0, 0 };
    private int drawCount;

    private int kernelID_frustumCull;
    private int WORKER_THREAD = 256;
    private Matrix4x4[] L2Ws_CPU;
    private InstanceBound[] worldBBoxBuffer_CPU;

    public ComputeShader CullingShader;
    private ComputeBuffer worldBBoxBuffer;
    private ComputeBuffer resultL2WBuffer;
    private ComputeBuffer allL2WBuffer;
    private ComputeBuffer bufferWithArgs;

    private int id_MAT_PV = Shader.PropertyToID("MAT_PV");
    private int id_instanceCount = Shader.PropertyToID("instanceCount");
    private int id_worldBBoxBuffer = Shader.PropertyToID("worldBBoxBuffer");
    private int id_useHiz = Shader.PropertyToID("useHizCull");
    private int id_resultL2WBuffer = Shader.PropertyToID("resultL2WBuffer");
    private int id_allL2WBuffer = Shader.PropertyToID("allL2WBuffer");
    private int id_bufferWithArgs = Shader.PropertyToID("bufferWithArgs");
    private int id_HiZBuffer = Shader.PropertyToID("HiZBuffer");

    void Start()
    {
        bufferWithArgs = new ComputeBuffer(1, args.Length * sizeof(uint), ComputeBufferType.IndirectArguments);
        UpdateBuffers();
        m_hizBuffer.InitializeHiZBuffer();
    }
    void OnPreCull()
    {
        if(useFrustumCull)
        {
            Profiler.BeginSample("CS: set culling params");
            Matrix4x4 mat_pv = mainCamera.projectionMatrix * mainCamera.worldToCameraMatrix;

            CullingShader.SetMatrix(id_MAT_PV, mat_pv);
            CullingShader.SetVector("textureSize", m_hizBuffer.textureSize);
            CullingShader.SetInt(id_instanceCount, (int)instanceCount);
            CullingShader.SetBuffer(kernelID_frustumCull, id_allL2WBuffer, allL2WBuffer);
            CullingShader.SetBuffer(kernelID_frustumCull, id_resultL2WBuffer, resultL2WBuffer);
            CullingShader.SetBuffer(kernelID_frustumCull, id_worldBBoxBuffer, worldBBoxBuffer);

            args[1] = 0;
            bufferWithArgs.SetData(args);

            CullingShader.SetBool(id_useHiz, useHiZCull);
            if (useHiZCull)
            {
                CullingShader.SetTexture(kernelID_frustumCull, id_HiZBuffer, m_hizBuffer.Texture);
            }
            CullingShader.SetBuffer(kernelID_frustumCull, id_bufferWithArgs, bufferWithArgs);
            CullingShader.Dispatch(kernelID_frustumCull, (int)instanceCount / WORKER_THREAD + 1, 1, 1);
            Profiler.EndSample();

            Profiler.BeginSample("CS : set material buffers");
            if(useFrustumCull)
            {
                instanceMaterial.SetBuffer("local2WorldBuffer", resultL2WBuffer);
            }
            else
            {
                instanceMaterial.SetBuffer("local2WorldBuffer", allL2WBuffer);
            }
            Profiler.EndSample();

            Profiler.BeginSample("get drawCount");
            //bufferWithArgs.GetData(args);
            //drawCount = (int)args[1];
            Profiler.EndSample();

            cachedInstanceCount = instanceCount;
            cachedSubMeshIndex = subMeshIndex;
        }

    }

    void Update()
    {
        if(useHiZCull)
        {
            m_hizBuffer.InitializeHiZBuffer();
        }
        // Update starting position buffer
        if (cachedInstanceCount != instanceCount || cachedSubMeshIndex != subMeshIndex)
        {
            UpdateBuffers();
        }

        if(useFrustumCull)
        {
            Graphics.DrawMeshInstancedIndirect(instanceMesh, subMeshIndex, instanceMaterial,
        new Bounds(Vector3.zero, new Vector3(100.0f, 100.0f, 100.0f)), bufferWithArgs);
        }
        else
        {
            drawCount = instanceCount;
            Draw();
        }
    }

    void OnGUI()
    {
        GUI.Label(new Rect(265, 25, 250, 30), "Instance Count: " + instanceCount.ToString() 
            + ", drawed instances : " + drawCount.ToString());
    }

    void UpdateBuffers()
    {
        ReleaseBuffers();
        // Ensure submesh index is in range
        if (instanceMesh != null)
            subMeshIndex = Mathf.Clamp(subMeshIndex, 0, instanceMesh.subMeshCount - 1);

        worldBBoxBuffer = new ComputeBuffer(instanceCount, 24);
        allL2WBuffer = new ComputeBuffer(instanceCount, 64);
        resultL2WBuffer = new ComputeBuffer(instanceCount, 64);

        L2Ws_CPU = new Matrix4x4[instanceCount];
        worldBBoxBuffer_CPU = new InstanceBound[instanceCount];

        Bounds bound = instanceMesh.bounds;

        for (int i = 0; i < instanceCount; i++)
        {
            float angle = Random.Range(0.0f, Mathf.PI * 2.0f);
            float distance = Random.Range(20.0f, 100.0f);
            float height = Random.Range(-10.0f, 10.0f);
            float size = Random.Range(0.05f, 0.8f);
            float r = Random.Range(0.0f, 360.0f);

            Vector2 data = rotate2D(Mathf.Sin(angle) * distance, Mathf.Cos(angle) * distance, size * size * 0.5f * r);
            Vector4 col0 = new Vector4(size, 0, 0, 0);
            Vector4 col1 = new Vector4(0, size, 0, 0);
            Vector4 col2 = new Vector4(0, 0, size, 0);
            Vector4 col3 = new Vector4(data.x, height, data.y, 1);

            L2Ws_CPU[i] = new Matrix4x4(col0, col1, col2, col3);
            if(useFrustumCull)
            {
                worldBBoxBuffer_CPU[i] = new InstanceBound();
                Vector4 center = L2Ws_CPU[i] * new Vector4(bound.center.x, bound.center.y, bound.center.z,1);
                worldBBoxBuffer_CPU[i].center  = new Vector3(center.x / center.w, center.y / center.w, center.z/ center.w);
                worldBBoxBuffer_CPU[i].extents = size * bound.extents;
            }
        }

        if (instanceMesh != null)
        {
            args[0] = (uint)instanceMesh.GetIndexCount(subMeshIndex);
            args[1] = (uint)instanceCount;
            args[2] = (uint)instanceMesh.GetIndexStart(subMeshIndex);
            args[3] = (uint)instanceMesh.GetBaseVertex(subMeshIndex);
        }
        else
        {
            args[0] = args[1] = args[2] = args[3] = 0;
        }
        
        allL2WBuffer.SetData(L2Ws_CPU);
        
        if (useFrustumCull)
        {
            bufferWithArgs.SetData(args);
            worldBBoxBuffer.SetData(worldBBoxBuffer_CPU);
        }
    }

    void Draw()
    {
        if (instanceMesh != null)
        {
            args[0] = (uint)instanceMesh.GetIndexCount(subMeshIndex);
            args[1] = (uint)instanceCount;
            args[2] = (uint)instanceMesh.GetIndexStart(subMeshIndex);
            args[3] = (uint)instanceMesh.GetBaseVertex(subMeshIndex);
        }
        else
        {
            args[0] = args[1] = args[2] = args[3] = 0;
        }
        bufferWithArgs.SetData(args);

        cachedInstanceCount = instanceCount;
        cachedSubMeshIndex = subMeshIndex;

        instanceMaterial.SetBuffer("local2WorldBuffer", allL2WBuffer);

        Graphics.DrawMeshInstancedIndirect(instanceMesh, subMeshIndex, instanceMaterial,
               new Bounds(Vector3.zero, new Vector3(100.0f, 100.0f, 100.0f)), bufferWithArgs);
    }

    void OnDisable()
    {
        ReleaseBuffers();
        if (bufferWithArgs != null)
            bufferWithArgs.Release();
        bufferWithArgs = null;
    }

    private Vector2 rotate2D(float x, float y, float r)
    {
        float s = Mathf.Sin(r);
        float c = Mathf.Cos(r);
        Vector2 ret = new Vector2(x * c - y * s, x * s + y * c);
        return ret;
    }
    private void ReleaseBuffers()
    {
        if (worldBBoxBuffer != null)
            worldBBoxBuffer.Release();
        if (allL2WBuffer != null)
            allL2WBuffer.Release();
        if (resultL2WBuffer != null)
            resultL2WBuffer.Release();
    }
}